'''
Process the historical daily data by adding indicators to the raw data.

Created on April, 2012
@author: bruce
'''

from datetime import date
from math import sqrt
from pprint import pprint

from dailyquote import DailyData

class HistdataBySymbol:
    def __init__(self, symbol, csv_dir="input\\", period=0):
        self._symbol = symbol
        self._histdata = [] 
        self._csv_path = csv_dir + symbol + ".csv"
        self._period = period
        self._size = 0                   

    """
    Get the historical daily raw data for the specific date range from the local data file.
    """
    def get_histdailyraws(self, start_date=None, end_date=None):  
        # Get the daily raw data from the csv file.      
        inp   = open(self._csv_path)
        lines = inp.readlines()          
        fields_str = lines[0].strip()
        fields = fields_str.split(',')
        
        def _get_datefromdaily(index):       
            # non-class helper function to get the date from quote line. 
            date_str = lines[index].split(',')[0]
            (y_str, m_str, d_str) = date_str.split('-')
            return date(int(y_str), int(m_str), int(d_str))   
                
        def _searchindex(date, start_index, end_index): 
            # non-class binary search helper function to find the index of the date. 
            index = 0
            # stop condition
            if start_index >= end_index:
                index = start_index
            else:
                mid_index = int((start_index + end_index) / 2)                
                if _get_datefromdaily(mid_index) < date:
                    index = _searchindex(date, mid_index + 1, end_index)
                else:
                    index = _searchindex(date, start_index, mid_index) 
            return index
        
        start_index = 1
        end_index = len(lines) - 1        
        if start_date != None and end_date != None:
            assert start_date <= end_date            
        if start_date != None:
            start_index = _searchindex(start_date, start_index, end_index)        
        if end_date != None:  
            end_index = _searchindex(end_date, start_index, end_index)        
        assert start_index <= end_index        
        quote_strs = lines[start_index : end_index + 1]             
        for quote_str in quote_strs:
            # Zip together the field names and values
            row = quote_str.strip().split(',')
            items = zip(fields, row)
            dailyRaw = {}
            # Add the value to our dictionary
            for (name, value) in items:
                dailyRaw[name] = value.strip()
            self._histdata.append(DailyData(dailyRaw))   
        self.size = len(self._histdata)                 
        
        
    """
    Caculate all the required indicators by taking arguments as indicator-period pairs.
    ind_pers: indicator - period pairs
    """
    def calculate_indicators(self, ind_pers):
        properties = ['o', 'c', 'h', 'lo', 'adj_c']                
        indicator_func = { 'ma':    self._calculate_ma,
                          'ema':    self._calculate_ema,
                          'std':    self._calculate_std,
                           'bb':    self._calculate_bb,
                          'wpr':    self._calculate_wpr,
                         'macd':    self._calculate_macd,
                          'env':    self._calculate_envelopes,
                          'rsi':    self._calculate_rsi,
                          'roc':    self._calculate_roc,
                           'mo':    self._calculate_momentum
                          }
                      
        for ind, period in ind_pers:
            if ind in properties:
                # No need to calculate properties
                continue
            assert ind in indicator_func.keys()
            indicator_func[ind](period)
                   
    """
    Get the quote of a specific date in the local data file.
    """                
    def get_quotebydate(self, dt, exact=True):            
        quote = self.__bSearchByDate(dt, 0, self.size - 1) 
        if exact and quote.dt != dt:
            return None
        else:    
            return quote
    
    """ 
    Use binary search to find the quote data of the date.
    """
    def __bSearchByDate(self, dt, start, end):
        if start > end:
            return self._histdata[end]
        mid = int((start + end) / 2)
        if self._histdata[mid].dt == dt:
            return self._histdata[mid]
        elif self._histdata[mid].dt > dt:
            return self.__bSearchByDate(dt, start, mid - 1)
        else:
            return self.__bSearchByDate(dt, mid + 1, end)
    
    """    
    Calculate moving average of the data as a dictionary of period.
    This method needs to be called after GetHistDailyRaws to be valid.
    Reference to http://en.wikipedia.org/wiki/Moving_average
    """
    def _calculate_ma(self, period):            
        self.__indicatorInitialize('ma')                
        if self.size > 0 and period not in self._histdata[0].ma.keys():
            for i in range(self.size):
                ave = 0
                period_len = min(period, i + 1)
                for n in range(i + 1 - period_len, i + 1):
                    ave += self._histdata[n].c
                ave = ave/period_len                    
                self._histdata[i].ma[period] = ave
     
    """           
    Calculates values of Exponential Moving Average for historic quote values as a dictionary of period.      
    ema is a dictionary with keys as period.
    Reference to http://en.wikipedia.org/wiki/Exponential_moving_average#Exponential_moving_average
    """
    def _calculate_ema(self, period):        
        self.__indicatorInitialize('ema')
        # If the period is already set in the dictionary, skip the calculation.
        if self.size > 0 and period not in self._histdata[0].ema.keys():
            self._histdata[0].ema[period] = self._histdata[0].c
            for i in range(1, self.size):
                exponent = 2.0 / (min(period, i + 1) + 1)
                self._histdata[i].ema[period] = (exponent * self._histdata[i].c + 
                                                 (1 - exponent) * self._histdata[i - 1].ema[period])                 
    
    """          
    Calculate the standard deviation of the data as a dictionary of period. 
    This method needs to call into _calculate_ma.   
    If we use population standard deviation, we need to add 1 to the denominator.
    Reference to http://en.wikipedia.org/wiki/Standard_deviation#Estimation
    """
    def _calculate_std(self, period, popstd=False):
        self.__indicatorInitialize('std')
        self._calculate_ma(period)
        
        if self.size > 0 and period not in self._histdata[0].std.keys():
            self._histdata[0].std[period] = 0
            for i in range(1, self.size):
                var = 0
                period_len = min(period, i + 1) # number of days from the data beginning
                for n in range(i + 1 - period_len, i + 1):
                    var += (self._histdata[n].c - self._histdata[i].ma[period])**2
                var /= period + popstd - 1 # True = 1
                    
                self._histdata[i].std[period] = sqrt(var)
    
    """    
    Calcuate Bollinger Bands Indicator of the data as a dictionary of period.
    This method needs to be called after _calculate_ma and CalculateStD to be valid.
    Reference to http://en.wikipedia.org/wiki/Bollinger_Bands
    """
    def _calculate_bb(self, period, k=1):
        self.__indicatorInitialize('bb')        
        if self.size > 0 and period not in self._histdata[0].bb.keys():    
            self._calculate_ma(period)
            self._calculate_std(period)        
            for i in range(1, self.size):             
                bb_l = self._histdata[i].ma[period] - self._histdata[i].std[period] * k
                bb_u = self._histdata[i].ma[period] + self._histdata[i].std[period] * k  
                self._histdata[i].bb[period] = (bb_l, bb_u)
    
    """
    Calculate values of Williams Percent Range for historic quote values as a dictionary of period.
    Reference to http://en.wikipedia.org/wiki/Williams_%25R
    %R= (Close(today) - High(Ndays)/(High(Ndays - Low(Ndays)) * 100, btw [-100, 0].
    """
    def _calculate_wpr(self, period):
        self.__indicatorInitialize('wpr')
        if self.size > 0 and period not in self._histdata[0].wpr.keys():
            self._histdata[0].wpr[period] = 50
            for i in range(1, self.size):
                period_len = min(i, period + 1)                        
                l = min([self._histdata[s].lo for s in range(i - period_len, i+1)])
                h = max([self._histdata[s].h for s in range(i - period_len, i+1)])
                # use the positive value in the code
                self._histdata[i].wpr[period] = (h - self._histdata[i].c) / (h - l) * 100 
            
    
    """
    Calculate values of moving average convergence/divergence for historic quote values.
    This method will call to _calculate_ema with two periods (fast & slow). 
    Reference to http://en.wikipedia.org/wiki/MACD
    """    
    def _calculate_macd(self, period=0):
        if self.size > 0 and 'macd' not in vars(self._histdata[0]):
            self._calculate_ema(12)
            self._calculate_ema(26)
            self._calculate_ema(9)
            for i in range(self.size):
                self._histdata[i].macd = self._histdata[i].ema[12] - self._histdata[i].ema[26]
    
    """
    Calculate values of envelopes for historic quote values.
    Envelope is defined as the bounds with buffer applied to the moving average: envelope = (mv - buffer, mv + buffer)
    It is different from the envelope described here: http://financial-dictionary.thefreedictionary.com/envelope.
    """
    def _calculate_envelopes(self, period, buff=0.10, is_fixed=False): # default to have 10% buffer
        self.__indicatorInitialize('env')
        if self.size > 0 and period not in self._histdata[0].env.keys():
            self._calculate_ma(period)
            for quote in self._histdata:
                if is_fixed:
                    quote.env[period] = (quote.ma[period] - buff, quote.ma[period] + buff)
                else:
                    quote.env[period] = (quote.ma[period] * (1 - buff), quote.ma[period] * (1 + buff))             
                        
    
    """
    Calculate values of relative strength index for historic quote values.
    We use rs = ma(u, n)/ma(d, n) here. Reference to http://en.wikipedia.org/wiki/Relative_Strength_Index 
    for an implementation with rs = ema(u, n)/ema(d, n)
    """
    def _calculate_rsi(self, period):
        self.__indicatorInitialize('rsi')
        if self.size > 0 and period not in self._histdata[0].rsi.keys():
            for i in range(self.size):
                period_len = min(i, period)
                up, down, upcount, downcount, aveup, avedown, rs = 0, 0, 0, 0, 0, 0, 0
                for s in range(i + 1 - period_len, i + 1):
                    # Note that the iteration goes from 1 to i + 1, since we need a previous value for comparison.
                    if self._histdata[s].c > self._histdata[s - 1].c:
                        up += self._histdata[s].c - self._histdata[s - 1].c
                        upcount += 1
                    elif self._histdata[s].c < self._histdata[s - 1].c:
                        down += self._histdata[s - 1].c - self._histdata[s].c
                        downcount += 1
                          
                if upcount > 0:
                    aveup = up/upcount
                if downcount > 0:
                    avedown = down/downcount
                if avedown != 0:
                    rs = aveup / avedown
                self._histdata[i].rsi[period] = 100 - 100/(1 + rs)
       
    """                       
    Calculate values of rate of change for historic quote values.
    roc = (close(today) - close(n days ago))/close(n days ago) * 100
    Reference to http://en.wikipedia.org/wiki/Rate_of_change_(technical_analysis)
    """    
    def _calculate_roc(self, period):
        self.__indicatorInitialize('roc')
        if self.size > 0 and period not in self._histdata[0].roc.keys():
            for i in range(self.size):
                period_len = min(i, period)
                self._histdata[i].roc[period] = ((self._histdata[i].c - self._histdata[i - period_len].c) / 
                                                 self._histdata[i - period_len].c * 100)
    
    """
    Calculate values of momentum for historic quote values.
    momentum = Close(today) / Close(n days ago) * 100% = 1 + ROC
    This method will call into calculate_roc.
    """
    def _calculate_momentum(self, period):        
        self.__indicatorInitialize('mo')
        if self.size > 0 and period not in self._histdata[0].mo.keys():
            self._calculate_roc(period)
            for i in range(self.size):
                self._histdata[i].mo[period] = self._histdata[i].roc[period] + 100
    

    
    """ 
    Initialize the indicator as a dictionary if it's not initialized yet. 
    """
    def __indicatorInitialize(self, indicator):
        if self.size > 0 and indicator not in vars(self._histdata[0]):
            for quote in self._histdata:
                vars(quote)[indicator] = {}

# Test of the class/methods        
if __name__ == "__main__":    
    symbol = "AAPL"
    _histdata = HistdataBySymbol(symbol)
    #_histdata.GetHistDailyRaws()
    _histdata.get_histdailyraws(date(2011, 2, 12), date(2012, 4, 12))
    pprint(symbol)    
    pprint(vars(_histdata._histdata[0]))
    pprint(vars(_histdata._histdata[1]))    
    
    _histdata._calculate_ma(30)    
    pprint(vars(_histdata._histdata[1]))    
    pprint(vars(_histdata._histdata[2]))    
    pprint(vars(_histdata._histdata[15]))
    
    _histdata._calculate_ema(50)
    pprint(vars(_histdata._histdata[1]))    
    pprint(vars(_histdata._histdata[2]))    
    pprint(vars(_histdata._histdata[15]))
        
    _histdata._calculate_std(50)
    pprint(vars(_histdata._histdata[1]))    
    pprint(vars(_histdata._histdata[2]))    
    pprint(vars(_histdata._histdata[15]))
        
            
    _histdata._calculate_bb(17)
    pprint(vars(_histdata._histdata[1]))    
    pprint(vars(_histdata._histdata[2]))    
    pprint(vars(_histdata._histdata[15]))
        
    _histdata._calculate_wpr(100)
    pprint(vars(_histdata._histdata[1]))    
    pprint(vars(_histdata._histdata[2]))    
    pprint(vars(_histdata._histdata[15]))
    
    _histdata._calculate_macd(200)
    pprint(vars(_histdata._histdata[30]))  
        
    _histdata._calculate_envelopes(13, 0.05)
    pprint(vars(_histdata._histdata[30]))    
    pprint(vars(_histdata._histdata[50]))    
    pprint(vars(_histdata._histdata[100]))
        
    _histdata._calculate_rsi(18)
    pprint(vars(_histdata._histdata[30]))    
    pprint(vars(_histdata._histdata[50]))    
    pprint(vars(_histdata._histdata[100]))
       
    _histdata._calculate_momentum(18)
    pprint(vars(_histdata._histdata[30]))    
    pprint(vars(_histdata._histdata[50]))    
    pprint(vars(_histdata._histdata[100]))
        
    _histdata._calculate_roc(18)
    pprint(vars(_histdata._histdata[30]))    
    pprint(vars(_histdata._histdata[50]))    
    pprint(vars(_histdata._histdata[100]))    
    
    _histdata.calculate_indicators((('ma', 20), ('roc', 30)))
    _histdata.calculate_indicators((('std', 40), ('rsi', 25)))
    pprint(vars(_histdata._histdata[0]))
    pprint(vars(_histdata._histdata[-1]))
    